"""
08_transaction_management.py
============================
Neo4j Transaction Management

This module covers:
1. ACID transactions
2. Read vs Write transactions
3. Transaction functions
4. Error handling and retries
5. Bookmarks for causal consistency

Author: Neo4j Tutorial
Date: January 2025
"""

from typing import Dict, List, Any, Callable, Optional
from contextlib import contextmanager
import time
import sys
import os

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from src.utils.connection import get_driver

from neo4j import Transaction, ManagedTransaction
from neo4j.exceptions import (
    TransientError,
    ServiceUnavailable,
    SessionExpired,
    ClientError
)


class TransactionManager:
    """
    Manages Neo4j transactions with proper patterns.
    
    Transaction Types:
    
    ┌─────────────────────────────────────────────────────────────────┐
    │                    TRANSACTION MODES                            │
    ├─────────────────────────────────────────────────────────────────┤
    │                                                                 │
    │  AUTO-COMMIT                                                    │
    │  ───────────                                                    │
    │  session.run(query)                                            │
    │  • Single query                                                │
    │  • Automatic commit                                            │
    │  • No explicit transaction control                             │
    │                                                                 │
    │  EXPLICIT TRANSACTION                                          │
    │  ────────────────────                                           │
    │  tx = session.begin_transaction()                              │
    │  tx.run(query1)                                                │
    │  tx.run(query2)                                                │
    │  tx.commit()  # or tx.rollback()                               │
    │  • Multiple queries                                            │
    │  • Manual commit/rollback                                      │
    │  • Full control                                                │
    │                                                                 │
    │  TRANSACTION FUNCTION (RECOMMENDED)                            │
    │  ────────────────────────────────────                           │
    │  session.execute_write(func)                                   │
    │  session.execute_read(func)                                    │
    │  • Automatic retries on transient errors                       │
    │  • Proper resource management                                  │
    │  • Best for production code                                    │
    │                                                                 │
    └─────────────────────────────────────────────────────────────────┘
    
    ACID Properties:
    - Atomicity: All operations succeed or all fail
    - Consistency: Database always in valid state
    - Isolation: Concurrent transactions don't interfere
    - Durability: Committed changes persist
    """
    
    def __init__(self, driver=None):
        self.driver = driver or get_driver()
        self._last_bookmarks = None
    
    # =========================================================
    # AUTO-COMMIT TRANSACTIONS
    # =========================================================
    
    def auto_commit_query(self, query: str, parameters: Dict = None) -> List[Dict]:
        """
        Execute a single query with auto-commit.
        
        Use for:
        - Simple read queries
        - Single operations that don't need grouping
        
        Note: Each run() is a separate transaction
        """
        with self.driver.session() as session:
            result = session.run(query, parameters or {})
            return [record.data() for record in result]
    
    # =========================================================
    # EXPLICIT TRANSACTIONS
    # =========================================================
    
    def explicit_transaction_example(self, patient_id: str, disease_id: str) -> bool:
        """
        Demonstrate explicit transaction with manual control.
        
        Use when you need:
        - Multiple operations that must succeed together
        - Manual commit/rollback control
        - Custom error handling between operations
        """
        with self.driver.session() as session:
            tx = session.begin_transaction()
            
            try:
                # Operation 1: Update patient
                tx.run(
                    "MATCH (p:Patient {id: $id}) SET p.lastVisit = datetime()",
                    id=patient_id
                )
                
                # Operation 2: Log the visit
                tx.run(
                    """
                    MATCH (p:Patient {id: $pid})
                    CREATE (v:Visit {
                        date: datetime(),
                        patientId: $pid,
                        type: 'checkup'
                    })
                    CREATE (p)-[:HAD_VISIT]->(v)
                    """,
                    pid=patient_id
                )
                
                # If we get here, commit
                tx.commit()
                return True
                
            except Exception as e:
                # Any error, rollback everything
                tx.rollback()
                print(f"Transaction failed: {e}")
                return False
    
    @contextmanager
    def transaction_context(self, database: str = None):
        """
        Context manager for explicit transactions.
        
        Usage:
            with manager.transaction_context() as tx:
                tx.run(query1)
                tx.run(query2)
            # Auto-commit on exit, rollback on exception
        """
        session = self.driver.session(database=database)
        tx = session.begin_transaction()
        
        try:
            yield tx
            tx.commit()
        except Exception:
            tx.rollback()
            raise
        finally:
            session.close()
    
    # =========================================================
    # TRANSACTION FUNCTIONS (RECOMMENDED)
    # =========================================================
    
    def execute_read(
        self,
        work: Callable[[ManagedTransaction], Any],
        **kwargs
    ) -> Any:
        """
        Execute a read transaction function.
        
        Benefits:
        - Automatic retries on transient errors
        - Proper connection management
        - Optimized for read operations
        """
        with self.driver.session() as session:
            return session.execute_read(work, **kwargs)
    
    def execute_write(
        self,
        work: Callable[[ManagedTransaction], Any],
        **kwargs
    ) -> Any:
        """
        Execute a write transaction function.
        
        Benefits:
        - Automatic retries on transient errors
        - Writes go to leader in cluster
        - ACID guarantees
        """
        with self.driver.session() as session:
            return session.execute_write(work, **kwargs)
    
    def create_patient_with_diagnosis(
        self,
        patient_data: Dict,
        disease_id: str,
        diagnosis_data: Dict
    ) -> Dict[str, Any]:
        """
        Example of transaction function pattern.
        
        All operations are atomic - either all succeed or all fail.
        """
        
        def _create_transaction(tx: ManagedTransaction) -> Dict:
            # Create patient
            result1 = tx.run(
                """
                MERGE (p:Patient {id: $id})
                SET p.name = $name,
                    p.createdAt = datetime()
                RETURN p
                """,
                **patient_data
            )
            patient = result1.single()["p"]
            
            # Create diagnosis relationship
            result2 = tx.run(
                """
                MATCH (p:Patient {id: $pid})
                MATCH (d:Disease {id: $did})
                MERGE (p)-[r:DIAGNOSED_WITH]->(d)
                SET r.diagnosedDate = date($date),
                    r.status = $status
                RETURN r
                """,
                pid=patient_data["id"],
                did=disease_id,
                date=diagnosis_data.get("date"),
                status=diagnosis_data.get("status", "active")
            )
            diagnosis = result2.single()["r"]
            
            return {
                "patient": dict(patient),
                "diagnosis": dict(diagnosis)
            }
        
        return self.execute_write(_create_transaction)
    
    # =========================================================
    # RETRY LOGIC
    # =========================================================
    
    def with_retry(
        self,
        operation: Callable,
        max_retries: int = 3,
        retry_delay: float = 1.0,
        **kwargs
    ) -> Any:
        """
        Execute operation with custom retry logic.
        
        Retries on:
        - TransientError (temporary issues)
        - ServiceUnavailable (connection issues)
        - SessionExpired (session timeout)
        """
        last_exception = None
        
        for attempt in range(max_retries + 1):
            try:
                return operation(**kwargs)
            
            except (TransientError, ServiceUnavailable, SessionExpired) as e:
                last_exception = e
                
                if attempt < max_retries:
                    wait_time = retry_delay * (2 ** attempt)  # Exponential backoff
                    print(f"Retry {attempt + 1}/{max_retries} after {wait_time}s: {e}")
                    time.sleep(wait_time)
                else:
                    print(f"Max retries reached: {e}")
            
            except ClientError as e:
                # Client errors should not be retried (syntax errors, etc.)
                raise
        
        raise last_exception
    
    # =========================================================
    # CAUSAL CONSISTENCY WITH BOOKMARKS
    # =========================================================
    
    def write_with_bookmark(
        self,
        work: Callable[[ManagedTransaction], Any],
        **kwargs
    ) -> Any:
        """
        Execute write and capture bookmark.
        
        Bookmarks ensure subsequent reads see the writes.
        Important in cluster deployments where reads may
        go to replicas.
        """
        with self.driver.session(
            bookmarks=self._last_bookmarks
        ) as session:
            result = session.execute_write(work, **kwargs)
            self._last_bookmarks = session.last_bookmarks()
            return result
    
    def read_with_bookmark(
        self,
        work: Callable[[ManagedTransaction], Any],
        **kwargs
    ) -> Any:
        """
        Execute read after previous writes.
        
        Uses bookmark to ensure read sees previous writes.
        """
        with self.driver.session(
            bookmarks=self._last_bookmarks
        ) as session:
            return session.execute_read(work, **kwargs)
    
    # =========================================================
    # BATCH OPERATIONS IN TRANSACTION
    # =========================================================
    
    def batch_create_in_transaction(
        self,
        items: List[Dict],
        batch_size: int = 1000
    ) -> int:
        """
        Batch create items within transactions.
        
        Uses UNWIND for efficiency while maintaining
        transaction boundaries.
        """
        total_created = 0
        
        def _create_batch(tx: ManagedTransaction, batch: List[Dict]) -> int:
            result = tx.run(
                """
                UNWIND $items AS item
                CREATE (n:BatchItem {
                    id: item.id,
                    name: item.name,
                    createdAt: datetime()
                })
                RETURN count(n) as created
                """,
                items=batch
            )
            return result.single()["created"]
        
        # Process in batches
        for i in range(0, len(items), batch_size):
            batch = items[i:i + batch_size]
            created = self.execute_write(_create_batch, batch=batch)
            total_created += created
        
        return total_created


# ============================================================
# USAGE EXAMPLES
# ============================================================

# Transaction function pattern (recommended)
def get_patient_summary(tx: ManagedTransaction, patient_id: str) -> Dict:
    """Example transaction function for reads."""
    result = tx.run(
        """
        MATCH (p:Patient {id: $id})
        OPTIONAL MATCH (p)-[:DIAGNOSED_WITH]->(d:Disease)
        OPTIONAL MATCH (p)-[:TAKES]->(m:Medication)
        RETURN p.name as name,
               collect(DISTINCT d.name) as diseases,
               collect(DISTINCT m.name) as medications
        """,
        id=patient_id
    )
    return result.single().data()


def update_patient_status(tx: ManagedTransaction, patient_id: str, active: bool) -> bool:
    """Example transaction function for writes."""
    result = tx.run(
        """
        MATCH (p:Patient {id: $id})
        SET p.active = $active, p.updatedAt = datetime()
        RETURN p.id as id
        """,
        id=patient_id,
        active=active
    )
    return result.single() is not None


# ============================================================
# DEMONSTRATION
# ============================================================

if __name__ == "__main__":
    print("=" * 60)
    print("Transaction Management - Demonstration")
    print("=" * 60)
    
    try:
        manager = TransactionManager()
        
        # Auto-commit example
        print("\n📌 Auto-commit Query")
        print("-" * 40)
        result = manager.auto_commit_query(
            "MATCH (p:Patient) RETURN p.name LIMIT 3"
        )
        print(f"   Found {len(result)} patients")
        
        # Transaction function example
        print("\n📌 Transaction Function (Recommended)")
        print("-" * 40)
        print("   Using execute_read for safe reads...")
        print("   Using execute_write for atomic writes...")
        
        # Context manager example
        print("\n📌 Transaction Context Manager")
        print("-" * 40)
        print("   with manager.transaction_context() as tx:")
        print("       tx.run(query1)")
        print("       tx.run(query2)")
        print("   # Auto-commit on success, rollback on error")
        
        print("\n" + "=" * 60)
        print("✅ Transaction management demonstration complete!")
        print("=" * 60)
        
    except Exception as e:
        print(f"Error: {e}")
